#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
menu_intraday_terminal_plus_v4.py
-------------------------------------------------
修复点（按你的建议全部采纳）：
1) 日线不再走 --tf，改为 --period daily，严格按 --start/--end 取数。
2) 分钟/小时自动根据 start~end 计算 --id-days（+缓冲），保证回放起点与输入一致。
3) 日线时间过滤放宽（不再强制 15:00），避免被误过滤。
4) “回放K线数量 bars” 参数真正生效，只播放指定根数。
5) 资产类型提示与兼容：将输入的 "fund" 映射为 "openfund"，提示改为 stock/etf/lof/openfund。
6) 其它：开盘价=0 回补为收盘价，统一 float 转换避免指标报错。
"""

import time
import subprocess
import pandas as pd
import pandas_ta as ta
from datetime import datetime, timedelta
import os

# --------------------- 小工具 ---------------------
def normalize_asset(asset: str) -> str:
    a = (asset or "").strip().lower()
    if a == "fund":
        return "openfund"
    return a

# --------------------- 数据获取 ---------------------
def fetch_data(symbol, asset, tf, start, end, session="cn", adjust="none"):
    """
    根据周期智能路由：
    - tf == '1d'  → 走日频：--period daily + --start/--end
    - 其它（1m/5m/30m/60m/2h/4h）→ 走 --tf 分支：自动计算 --id-days 覆盖 start~end 区间
    """
    tmpfile = "temp_fetch.csv"
    start_dt = pd.to_datetime(start).date()
    end_dt = pd.to_datetime(end).date()

    cmd = [
        "python3", "ak_a_hist_plus_fund.py",
        "--asset", asset,
        "--symbol", symbol,
    ]

    if tf == "1d":
        # 日线：严格使用 start/end
        cmd += [
            "--period", "daily",
            "--start", start_dt.strftime("%Y-%m-%d"),
            "--end", end_dt.strftime("%Y-%m-%d"),
            "--adjust", adjust,
            "--export", tmpfile,
        ]
    else:
        # 分钟/小时：底层会忽略 start/end，按 id-days 回看分钟数据再重采样。
        days = (end_dt - start_dt).days
        id_days = max(2, days + 2)  # +2 天缓冲，防止节假日/停牌导致覆盖不足
        cmd += [
            "--start", start_dt.strftime("%Y-%m-%d"),  # 传了也没关系，--tf 路径里会以 id-days 为准
            "--end", end_dt.strftime("%Y-%m-%d"),
            "--tf", tf,
            "--id-days", str(id_days),
            "--session", session,
            "--adjust", adjust,
            "--export", tmpfile,
        ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            return pd.DataFrame()
        if not os.path.exists(tmpfile):
            return pd.DataFrame()
        df = pd.read_csv(tmpfile)
        os.remove(tmpfile)
        column_mapping = {
            "日期": "time",
            "时间": "time",
            "开盘": "open",
            "开盘价": "open",
            "最高": "high",
            "最高价": "high",
            "最低": "low",
            "最低价": "low",
            "收盘": "close",
            "收盘价": "close",
            "成交量": "volume",
            "成交额": "amount",
        }
        df = df.rename(columns=column_mapping)
        df["time"] = pd.to_datetime(df["time"])
        df = df.set_index("time")
        for col in ["open", "high", "low", "close", "volume", "amount"]:
            df[col] = pd.to_numeric(df[col], errors="coerce").fillna(0.0)
        # 开盘价为 0 时用收盘价回补
        df.loc[df["open"] == 0, "open"] = df["close"]
        return df
    except Exception as e:
        print("❌ 数据获取错误:", e)
        return pd.DataFrame()

# --------------------- 时间控制 ---------------------
def get_cn_trading_sessions(date=None):
    if date is None:
        date = datetime.now().date()
    morning_start = datetime.combine(date, datetime.strptime("09:30", "%H:%M").time())
    morning_end = datetime.combine(date, datetime.strptime("11:30", "%H:%M").time())
    afternoon_start = datetime.combine(date, datetime.strptime("13:00", "%H:%M").time())
    afternoon_end = datetime.combine(date, datetime.strptime("15:00", "%H:%M").time())
    return [(morning_start, morning_end), (afternoon_start, afternoon_end)]

def is_trading_time(now=None):
    if now is None:
        now = datetime.now()
    sessions = get_cn_trading_sessions(now.date())
    for start, end in sessions:
        if start <= now <= end:
            return True
    return False

def next_open_time(now=None):
    if now is None:
        now = datetime.now()
    sessions = get_cn_trading_sessions(now.date())
    for start, end in sessions:
        if now < start:
            return start
    next_day = now.date() + timedelta(days=1)
    return datetime.combine(next_day, datetime.strptime("09:30", "%H:%M").time())

def valid_cn_time(ts, tf):
    # 放宽日线：CSV 常见为 00:00:00 或 15:00:00，均放行
    if tf == "1d":
        return True
    hm = ts.strftime("%H:%M")
    if tf in ["1m", "5m"]:
        return True
    elif tf == "30m":
        return hm in ["10:00", "10:30", "11:00", "11:30", "13:30", "14:00", "14:30", "15:00"]
    elif tf == "60m":
        return hm in ["10:30", "11:30", "14:00", "15:00"]
    return True

def is_recent_signal(ts, tf):
    now = datetime.now()
    delta = now - ts
    if tf == "30m":
        return delta <= timedelta(minutes=30)
    elif tf == "60m":
        return delta <= timedelta(minutes=60)
    elif tf == "1d":
        return delta <= timedelta(hours=24)
    return True

# --------------------- 指标检测 ---------------------
def detect_signals(df: pd.DataFrame):
    """
    只返回“触发型”信号：
    - KDJ：金叉/死叉；跌入超卖(<20)/升破超买(>80)
    - MACD：金叉/死叉；上穿零轴/下穿零轴
    - RSI：跌入超卖(<20)/升入超买(>80)
    注意：不返回“弱势/强势/放量/缩量”这类持续状态，避免每根都触发。
    """
    signals = []
    if df.empty or len(df) < 35:
        return signals

    df = df.copy().astype({
        "open": "float", "high": "float", "low": "float",
        "close": "float", "volume": "float", "amount": "float"
    })

    # --- KDJ ---
    kdj_df = ta.kdj(high=df["high"], low=df["low"], close=df["close"], length=9, signal=3)
    if isinstance(kdj_df, pd.DataFrame):
        cols_upper = [c.upper() for c in kdj_df.columns]
        def pick(letter):
            for c, cu in zip(kdj_df.columns, cols_upper):
                if cu.startswith(letter):
                    return c
            raise KeyError(f"KDJ列缺失: 需要 {letter}* in {kdj_df.columns.tolist()}")
        df["K"] = pd.to_numeric(kdj_df[pick("K")], errors="coerce")
        df["D"] = pd.to_numeric(kdj_df[pick("D")], errors="coerce")
        df["J"] = pd.to_numeric(kdj_df[pick("J")], errors="coerce")
    else:
        k, d, j = kdj_df
        df["K"] = pd.to_numeric(k, errors="coerce")
        df["D"] = pd.to_numeric(d, errors="coerce")
        df["J"] = pd.to_numeric(j, errors="coerce")

    # --- MACD ---
    macd = ta.macd(df["close"], fast=12, slow=26, signal=9)
    df["DIF"] = pd.to_numeric(macd["MACD_12_26_9"], errors="coerce")
    df["DEA"] = pd.to_numeric(macd["MACDs_12_26_9"], errors="coerce")
    # df["MACD_H"] = pd.to_numeric(macd["MACDh_12_26_9"], errors="coerce")  # 这次不用它做触发

    # --- RSI ---
    df["RSI6"]  = pd.to_numeric(ta.rsi(df["close"], length=6),  errors="coerce")
    df["RSI12"] = pd.to_numeric(ta.rsi(df["close"], length=12), errors="coerce")
    df["RSI24"] = pd.to_numeric(ta.rsi(df["close"], length=24), errors="coerce")

    last, prev = df.iloc[-1], df.iloc[-2]

    # KDJ 金叉/死叉
    if pd.notna(prev["K"]) and pd.notna(prev["D"]) and pd.notna(last["K"]) and pd.notna(last["D"]):
        if prev["K"] <= prev["D"] and last["K"] > last["D"]:
            signals.append("KDJ金叉")
        if prev["K"] >= prev["D"] and last["K"] < last["D"]:
            signals.append("KDJ死叉")

    # KDJ 阈值触发（进入区间的瞬间才触发）
    if pd.notna(prev["K"]) and pd.notna(last["K"]):
        if prev["K"] >= 80 and last["K"] < 80:
            signals.append("KDJ跌破超买")   # 看空触发
        if prev["K"] <= 20 and last["K"] > 20:
            signals.append("KDJ脱离超卖")   # 看多触发
        if prev["K"] >= 20 and last["K"] < 20:
            signals.append("KDJ跌入超卖")   # 有人喜欢把“跌入超卖”当博弈多头触发

    # MACD 金叉/死叉
    if pd.notna(prev["DIF"]) and pd.notna(prev["DEA"]) and pd.notna(last["DIF"]) and pd.notna(last["DEA"]):
        if prev["DIF"] <= prev["DEA"] and last["DIF"] > last["DEA"]:
            signals.append("MACD金叉")
        if prev["DIF"] >= prev["DEA"] and last["DIF"] < last["DEA"]:
            signals.append("MACD死叉")

    # MACD 零轴穿越
    if pd.notna(prev["DIF"]) and pd.notna(last["DIF"]):
        if prev["DIF"] <= 0 and last["DIF"] > 0:
            signals.append("MACD上穿零轴")
        if prev["DIF"] >= 0 and last["DIF"] < 0:
            signals.append("MACD下穿零轴")

    # RSI 阈值触发（进入区间的瞬间）
    if pd.notna(prev["RSI6"]) and pd.notna(last["RSI6"]):
        if prev["RSI6"] >= 80 and last["RSI6"] < 80:
            signals.append("RSI跌破超买")   # 看空触发
        if prev["RSI6"] <= 20 and last["RSI6"] > 20:
            signals.append("RSI脱离超卖")   # 看多触发
        if prev["RSI6"] >= 20 and last["RSI6"] < 20:
            signals.append("RSI跌入超卖")   # 有人当作抄底触发

    return signals



# --------------------- 信号输出 ---------------------
TF_STRENGTH = {
    "1m": "参考性较弱（噪声大）",
    "5m": "参考性一般",
    "30m": "参考性较强",
    "60m": "参考性最强",
    "1d": "长周期强参考",
}
TIMEFRAMES = ["1m", "5m", "30m", "60m"]

def print_kline(df: pd.DataFrame, tf: str, symbol: str):
    if df.empty:
        return
    row = df.iloc[-1]
    if not valid_cn_time(row.name, tf):
        return
    print(
        f"[{row.name}] {symbol} {tf} K线: "
        f"开盘={row['open']:.2f} 最高={row['high']:.2f} 最低={row['low']:.2f} 收盘={row['close']:.2f} "
        f"成交量={row['volume']} 成交额={row['amount']}"
    )

def print_signal(row: dict, strong=False):
    tf_strength = TF_STRENGTH.get(row["tf"], "参考性未知")
    color = "\x1b[92m" if row["direction"] == "long" else "\x1b[91m"
    reset = "\x1b[0m"
    border = "=" * 60 if strong else "-" * 60
    print(color + " " + border)
    title = "⚡ 强交易信号" if strong else "⚠️ 参考信号"
    print(f"{title} ({row['direction'].upper()})")
    print(f"时间: {row['time']}")
    print(f"标的: {row['symbol']}  周期: {row['tf']} ({tf_strength})")
    print(f"入场价: {row['entry']:.2f}")
    print(f"止损价: {row['stop']:.2f}")
    print(f"目标1: {row['tp1']:.2f}  目标2: {row['tp2']:.2f}")
    print(f"风险金额: {row['risk_amt']:.2f}  建议仓位: {row['size_suggest']} 手")
    print(f"信号原因: {' + '.join(row['signals'])}")
    print(border + reset + "\n")

# --------------------- 历史回放 ---------------------
def simulate_future(df, start_idx, hold_n=3):
    sub = df.iloc[start_idx:start_idx + hold_n]
    if sub.empty:
        return (0, 0, 0)
    entry = df.iloc[start_idx]["close"]
    max_up = (sub["high"].max() - entry) / entry * 100
    max_down = (sub["low"].min() - entry) / entry * 100
    final = (sub.iloc[-1]["close"] - entry) / entry * 100
    return (max_up, max_down, final)

def history_query(symbol, asset):
    tf = input("选择周期 (1m/5m/30m/60m/1d): ").strip()
    start = input("输入开始日期 (YYYY-MM-DD): ").strip()
    bars = int(input("输入回放K线数量: ").strip() or 50)

    df = fetch_data(symbol, asset, tf, start, datetime.now().strftime("%Y-%m-%d"))
    if df.empty:
        print("❌ 无法获取数据")
        return

    start_dt = pd.to_datetime(start)
    df = df[df.index >= start_dt]

    # 只回放 bars 根（从开始日期起向后）
    df_iter = df.iloc[:bars]

    for i in range(len(df_iter)):
        row = df_iter.iloc[i]
        if not valid_cn_time(row.name, tf):
            continue
        print(
            f"[{row.name}] {symbol} {tf} K线: 开盘={row['open']:.2f} 最高={row['high']:.2f} "
            f"最低={row['low']:.2f} 收盘={row['close']:.2f} 成交量={row['volume']} 成交额={row['amount']}"
        )
        signals = detect_signals(df_iter.iloc[: i + 1])
        if signals:
            # direction = "long" if any(
            #     "金叉" in s or "超卖" in s or "强势" in s for s in signals
            # ) else "short"
            bullish_keys = ("金叉", "脱离超卖", "跌入超卖", "上穿零轴")  # 你也可以只保留前两种更保守
            direction = "long" if any(any(k in s for k in bullish_keys) for s in signals) else "short"

            stop = row["low"] if direction == "long" else row["high"]
            tp1 = row["close"] * (1.005 if direction == "long" else 0.995)
            tp2 = row["close"] * (1.01 if direction == "long" else 0.99)
            risk_amt = 100000 * 0.005
            signal = {
                "time": row.name,
                "symbol": symbol,
                "tf": tf,
                "direction": direction,
                "entry": row["close"],
                "stop": stop,
                "tp1": tp1,
                "tp2": tp2,
                "risk_amt": risk_amt,
                "size_suggest": int(risk_amt / (abs(row["close"] - stop) or 1)),
                "signals": signals,
            }
            strong = len(signals) >= 2
            print_signal(signal, strong=strong)
            max_up, max_down, final = simulate_future(df_iter, i, hold_n=3)
            print(
                f"📊 模拟结果 (持有3根K线): 最大浮盈: {max_up:+.2f}% 最大浮亏: {max_down:+.2f}% 最终盈亏: {final:+.2f}%"
            )
            print("-" * 60)
        cmd = input("(回车继续，输入 q 退出): ").strip().lower()
        if cmd == "q":
            break

# --------------------- 实盘 ---------------------
def realtime_mode(symbol, asset, account, risk_pct, show_1m, show_expired):
    printed_times = set()
    try:
        while True:
            now = datetime.now()
            if not is_trading_time(now):
                nxt = next_open_time(now)
                print(
                    f"[未开盘] 当前时间：{now.strftime('%Y-%m-%d %H:%M:%S')}  下一个开盘时间：{nxt.strftime('%Y-%m-%d %H:%M:%S')}"
                )
                time.sleep(60)
                continue
            for tf in TIMEFRAMES:
                start = (now - timedelta(days=3)).strftime("%Y-%m-%d")
                end = now.strftime("%Y-%m-%d")
                df = fetch_data(symbol, asset, tf, start, end)
                if df.empty:
                    continue
                row = df.iloc[-1]
                if not valid_cn_time(row.name, tf):
                    continue
                print_kline(df, tf, symbol)
                signals = detect_signals(df)
                if not signals:
                    continue
                strong = len(signals) >= 2
                if tf == "1m" and not show_1m and not strong:
                    continue
                # direction = "long" if any(
                #     "金叉" in s or "超卖" in s or "强势" in s for s in signals
                # ) else "short"
                bullish_keys = ("金叉", "脱离超卖", "跌入超卖", "上穿零轴")  # 你也可以只保留前两种更保守
                direction = "long" if any(any(k in s for k in bullish_keys) for s in signals) else "short"

                stop = row["low"] if direction == "long" else row["high"]
                tp1 = row["close"] * (1.005 if direction == "long" else 0.995)
                tp2 = row["close"] * (1.01 if direction == "long" else 0.99)
                signal = {
                    "time": row.name,
                    "symbol": symbol,
                    "tf": tf,
                    "direction": direction,
                    "entry": row["close"],
                    "stop": stop,
                    "tp1": tp1,
                    "tp2": tp2,
                    "risk_amt": account * risk_pct,
                    "size_suggest": int((account * risk_pct) / (abs(row["close"] - stop) or 1)),
                    "signals": signals,
                }
                key = (signal["time"], tf)
                if key not in printed_times:
                    if not show_expired and not is_recent_signal(row.name, tf):
                        continue
                    printed_times.add(key)
                    print_signal(signal, strong=strong)
            time.sleep(30)
    except KeyboardInterrupt:
        print("\n用户手动停止运行。")

# --------------------- 主入口 ---------------------
def main():
    print("请选择功能：")
    print("1. 实盘分析（实时监控）")
    print("2. 历史K线查询（含信号和模拟盈亏）")
    choice = input("> ").strip()
    if choice == "1":
        symbol = input("请输入股票/基金代码: ").strip()
        asset = normalize_asset(input("请输入资产类型 (stock/etf/lof/openfund): ").strip())
        account = float(input("请输入账户金额: ").strip() or 100000)
        risk_pct = float(input("请输入风险比例(默认0.005): ").strip() or 0.005)
        show_1m = input("是否显示 1m 信号？(y/n): ").strip().lower() == "y"
        show_expired = input("是否显示过期信号？(y/n): ").strip().lower() == "y"
        realtime_mode(symbol, asset, account, risk_pct, show_1m, show_expired)
    elif choice == "2":
        symbol = input("请输入股票/基金代码: ").strip()
        asset = normalize_asset(input("请输入资产类型 (stock/etf/lof/openfund): ").strip())
        history_query(symbol, asset)
    else:
        print("无效选择")

if __name__ == "__main__":
    main()

